"""
邢不行 | 2023分享会
author: 邢不行
微信: xbx6660
"""
import pandas as pd
import warnings
from program.function import draw_hedge_signal_plotly, get_trade_info
import os

warnings.filterwarnings('ignore')

_ = os.path.abspath(os.path.dirname(__file__))  # 返回当前文件路径
root_path = os.path.abspath(os.path.join(_, '../'))  # 返回根目录文件夹

# =====需要配置的参数
start = '2023-02-21'  # 分析开始时间
end = '2023-07-10'  # 分析结束时间
period = '1H'  # 策略持仓周期，填入'1D'、'4H'这样的时间信息，或者填入auto自动计算策略的持仓周期（有过滤慎用auto）
name = 'MTMMEAN'  # 因子名称
k_path = '/Users/winkey/Documents/量化/2023分享会保温杯中性策略2期/2023分享会保温杯中性策略回测_v2/data/k线数据/'  # K线储存的路径
select_path = '//data/input/select.csv'  # 选币数据的路径

# =====数据的处理及周期处理
# ===读取选币数据 & 截取数据
# 读取选币数据
select = pd.read_csv(select_path, encoding='gbk', parse_dates=['candle_begin_time'],index_col=0)
print(select)
select = select[['candle_begin_time','ret_next','offset','symbol','方向']]
print(select)
# print(select.columns)

# 只保留分析区间内的选币数据
select = select[select['candle_begin_time'] >= pd.to_datetime(start)]
select = select[select['candle_begin_time'] <= pd.to_datetime(end)]

# ===计算period
if period == 'auto':
    # 如果是auto模式，自动计算第一个周期和第二个周期的时间差作为持仓周期
    date_list = sorted(select['candle_begin_time'].unique())  # 获取所有换仓日期
    period = pd.to_timedelta(date_list[1] - date_list[0])  # 第二次换仓日期 - 第一次换仓日期
else:
    period = pd.to_timedelta(period)  # 非auto模式，直接按照给的参数计算周期

# 计算一天有多少个period，为之后算日化收益做准备
daily_periods = pd.to_timedelta('1D') / period
# 计算这个区间内有多个个period，为之后画图做准备
total_periods = (pd.to_datetime(end) - pd.to_datetime(start)) / period

# =====按照币种及方向进行分组，计算一些分析所需要的数据
groups = select.groupby(['symbol', '方向'])
res_list = []  # 储存分组结果的list
# 遍历各个分组
for t, g in groups:
    # ===分组处理数据
    g.sort_values(by='candle_begin_time', inplace=True)
    # 计算前后K线的时间间隔
    g['span'] = g['candle_begin_time'].diff()
    # 时间间隔不等于设定周期的肯定是中间有过换仓，标记一下换仓的开始时间 (修改的时候可以把它改为'1h)
    g.loc[g['span'] != period, 'start_time'] = g['candle_begin_time']
    g['start_time'].fillna(method='ffill', inplace=True)
    # 按照不同的持仓时间段分类，并计算持仓列表
    hold = g.groupby('start_time').agg({'candle_begin_time': 'last'}).reset_index()
    hold['hold_info'] = hold['start_time'].apply(str) + '--' + hold['candle_begin_time'].apply(str)
    # 结合持仓数据，计算真实的收益（忽略手续费）
    g['ret_next_real'] = g['ret_next'] * t[1]

    # ===统计结果
    res = pd.DataFrame()  # 需要返回的结果
    res.loc[0, 'symbol'] = t[0]  # 币种名称
    res.loc[0, '方向'] = t[1]  # 交易方向
    res.loc[0, '选中次数'] = len(g['start_time'].unique())  # 区间内选中币种的次数
    res.loc[0, '累计周期数'] = g.shape[0]  # 区间内币种的持仓周期数
    res.loc[0, '每周期平均收益（日化）'] = (g['ret_next_real'].mean() + 1) ** daily_periods - 1  # 不同周期的平均日化收益
    # 区间内，币种的累计收益（准确的说是净值的百分比）
    if t[1] > 0:
        res.loc[0, '区间累计收益'] = (g['ret_next'] + 1).prod() - 1  # 做多的最终收益
    else:
        res.loc[0, '区间累计收益'] = 1 - (g['ret_next'] + 1).prod()  # 做空的最终收益
    res.loc[0, '首次选中时间'] = sorted(g['start_time'].unique())[0]  # 区间内首次选中币种的时间
    res.loc[0, '最后选中时间'] = sorted(g['start_time'].unique())[-1]  # 区间内最后一次选中币种的时间

    # 计算区间内的的持币时间段，格式如：['2021-05-09 13:00:00--2021-05-09 15:00:00', '2021-05-10 19:00:00--2021-05-11 05:00:00']
    res['选中周期'] = ''  # 小tips：需要往DataFrame的cell里面插入list，这一列需要是object类型（所以这里给了''，字符串就是object）
    res.at[0, '选中周期'] = hold['hold_info'].to_list()  # 往数据中插入list时，需要用at函数，loc不行。

    # 将计算的结果添加到结果汇总中
    res_list.append(res)

# =====汇总所有结果，再按照方向拆成多头和空头
# 汇总所有分组的分析结果
all_res = pd.concat(res_list, ignore_index=True)
# 拆分多头的分析结果
long_res = all_res[all_res['方向'] > 0].reset_index(drop=True)
# 拆分空头的分析结果
short_res = all_res[all_res['方向'] < 0 ].reset_index(drop=True)

# =====针对多空分析结果进行进一步分析
describe = pd.DataFrame()  # 分析结果储存的df

# 1.分析多头数据
describe.loc['选币数', '多头'] = long_res.shape[0]
describe.loc['平均选中次数', '多头'] = long_res['选中次数'].mean()
describe.loc['平均周期数', '多头'] = long_res['累计周期数'].mean()
describe.loc['平均日化收益', '多头'] = long_res['每周期平均收益（日化）'].mean()
describe.loc['平均累计收益', '多头'] = long_res['区间累计收益'].mean()
describe.loc['胜率_日均', '多头'] = long_res[long_res['每周期平均收益（日化）'] > 0].shape[0] / long_res.shape[0]
describe.loc['胜率_累计', '多头'] = long_res[long_res['区间累计收益'] > 0].shape[0] / long_res.shape[0]

# 2.分析空头数据
describe.loc['选币数', '空头'] = short_res.shape[0]
describe.loc['平均选中次数', '空头'] = short_res['选中次数'].mean()
describe.loc['平均周期数', '空头'] = short_res['累计周期数'].mean()
describe.loc['平均日化收益', '空头'] = short_res['每周期平均收益（日化）'].mean()
describe.loc['平均累计收益', '空头'] = short_res['区间累计收益'].mean()
describe.loc['胜率_日均', '空头'] = short_res[short_res['每周期平均收益（日化）'] > 0].shape[0] / short_res.shape[0]
describe.loc['胜率_累计', '空头'] = short_res[short_res['区间累计收益'] > 0].shape[0] / short_res.shape[0]
# 打印分析结果
print(describe)

# =====结果保存
# 保存数据的文件夹是否存在
save_path = root_path + f'/data/多空分析/{name}_{start}_{end}/'
if not os.path.exists(save_path):
    os.mkdir(save_path)
# 保存多头、空头的分析结果数据 及 汇总数据
long_res.to_csv(save_path + '01_多头分析结果.csv', encoding='gbk', index=False)
short_res.to_csv(save_path + '02_空头分析结果.csv', encoding='gbk', index=False)
describe.to_csv(save_path + '03_分析汇总.csv', encoding='gbk')

# =====根据多空结果，进行绘图
# 绘制的时候K线向历史和前后多扩展15%（最少扩展20个周期）
# K线开始时间
k_start = pd.to_datetime(start) - pd.to_timedelta(period) * max(int(total_periods * 0.15), 20)  # 周期数据开始时间
d_start = pd.to_datetime(start) - pd.to_timedelta('90D')  # 日线数据开始时间
# K线结束时间
k_end = pd.to_datetime(end) + pd.to_timedelta(period) * max(int(total_periods * 0.15), 20)  # 周期数据结束时间
d_end = pd.to_datetime(end) + pd.to_timedelta('90D')  # 日线数据结束时间
# 画图需要的信息
draw_info = {'多头': long_res, '空头': short_res}

# 开始绘制多空持仓的图片
for direction, info_df in draw_info.items():
    # 判断开仓信号是做多还是做空
    open_signal = 1 if direction == '多头' else -1
    fig_save_path = save_path + f'/{direction}行情图/'
    if not os.path.exists(fig_save_path):
        os.mkdir(fig_save_path)
    # 开始遍历每一行数据画图
    for i in info_df.index:
        # 获取币种名称
        coin_symbol = info_df.loc[i, 'symbol']
        print(f'正在绘制：{direction}_{coin_symbol}')
        coin_symbol = coin_symbol[:-4] + '-' + coin_symbol[-4:]
        # 读取币种信息
        df = pd.read_csv(k_path + coin_symbol + '.csv', encoding='gbk', skiprows=1,
                         parse_dates=['candle_begin_time'])
        # ===按照周期生成数据
        # 截取区间内的币种信息
        symbol_df = df[df['candle_begin_time'] >= pd.to_datetime(k_start)]
        symbol_df = symbol_df[symbol_df['candle_begin_time'] <= pd.to_datetime(k_end)]

        # 获取所有的开仓时间点
        open_times = [pd.to_datetime(time_range.split('--')[0]) for time_range in info_df.loc[i, '选中周期']]
        # 获取所有的平仓时间点
        close_times = [pd.to_datetime(time_range.split('--')[1]) for time_range in info_df.loc[i, '选中周期']]

        # 在数据中加入开仓信息
        symbol_df.loc[symbol_df['candle_begin_time'].isin(open_times), 'open_signal'] = open_signal
        # 在数据中加入平仓信息
        symbol_df.loc[symbol_df['candle_begin_time'].isin(close_times), 'close_signal'] = 0

        # 根据symbol_df生成每笔交易信息
        trade_df = get_trade_info(symbol_df, direction)

        # ===按照日线生成数据
        # 截取区间内的币种信息
        day_df = df[df['candle_begin_time'] >= pd.to_datetime(d_start)]
        day_df = day_df[day_df['candle_begin_time'] <= pd.to_datetime(d_end)]
        # 数据resample时的规则
        rule_dict = {'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last'}
        day_df = day_df.resample(rule='1D', on='candle_begin_time').agg(rule_dict).reset_index()

        # 在日线上标记区间开始时间
        if day_df['candle_begin_time'].dt.date.to_list()[0] < pd.to_datetime(start).date():
            day_df.loc[day_df['candle_begin_time'].dt.date == pd.to_datetime(start).date(), 'signal'] = 'start'
        else:
            day_df.loc[0, 'signal'] = 'start'

        # 在日线上标记区间结束时间
        if day_df['candle_begin_time'].dt.date.to_list()[-1] > pd.to_datetime(end).date():
            day_df.loc[day_df['candle_begin_time'].dt.date == pd.to_datetime(end).date(), 'signal'] = 'end'
        else:
            day_df.loc[day_df.index.max(), 'signal'] = 'end'

        # 绘制中性策略的开平仓信息
        draw_hedge_signal_plotly(symbol_df, fig_save_path, f'{direction}_{coin_symbol}', day_df=day_df,
                                 trade_df=trade_df)
